package rbtree;

import com.sun.org.apache.regexp.internal.RE;

import static rbtree.RBTreeNode.BLACK;
import static rbtree.RBTreeNode.RED;

public class RBTree {

    // 保存红黑树的根节点
    public RBTreeNode root = null;

    // 保存红黑树中的结点个数
    public int size = 0;

    /**
     * 向红黑树中插入新的关键字
     * return 是否插入成功
     * true : 插入成功
     * false : 插入失败 (出现重复)
     */
    public boolean insert(long key) {
        // 1. 按照普通搜索树的规则插入
        if (root == null) {
            root = new RBTreeNode(key, null, BLACK);
            size++;
            return true;
        }

        RBTreeNode parent = null;
        RBTreeNode current = root;

        while (current != null) {
            if (key == current.key) {
                return false;
            } else if (key < current.key) {
                parent = current;
                current = current.left;
            } else {
                parent = current;
                current = current.right;
            }
        }

        // 根据插入规则 每次插入的新结点一定是红色的
        RBTreeNode node = new RBTreeNode(key, parent, RED);
        if (key < parent.key) {
            parent.left = node;
        } else {
            parent.right = node;
        }
        size++;

        // 进行红黑树规则的判断 + 平衡调整过程
        adjustBalance(node);

        return true;
    }

    private void adjustBalance(RBTreeNode node) {
        while (true) {
            RBTreeNode parent = node.parent;
            if (parent == null) {
                break;
            }
            if (parent.color = BLACK) {
                break;
            }

            // 一定破坏了 "红色不能相邻" 的规则
            RBTreeNode grandpa = parent.parent;

            // 找到叔叔结点
            if (parent == grandpa.left) {
                RBTreeNode uncle = grandpa.right;
                if (uncle != null && uncle.color == RED) {
                    /**
                     * 情况一 : 叔叔存在 并且颜色是红色
                     * 步骤:
                     *      1. 叔叔和父亲的颜色改为黑色
                     *      2. 祖父的颜色改为红色
                     *      3. 把祖父视为 node, 再去判断是否违反规则了
                     */
                    parent.color = uncle.color = BLACK;
                    grandpa.color = RED;
                    node = grandpa;
                    continue;
                } else {
                    // 判断 grandpa <-> parent 和 parent <-> node 的关系是否不一致
                    // 已知 parent 是 grandpa 的左边
                    if (node == parent.right) {
                        leftRotate(parent);

                        // swap(parent,node);
                        parent = node;
                    }
                    // 接下来统一处理关系一致的情况
                    rightRotate(grandpa);
                    grandpa.color = RED;
                    parent.color = BLACK;
                    break;
                }
            } else {
                RBTreeNode uncle = grandpa.left;
                if (uncle != null && uncle.color == RED) {
                    /**
                     * 情况一 : 叔叔存在 并且颜色是红色
                     * 步骤:
                     *      1. 叔叔和父亲的颜色改为黑色
                     *      2. 祖父的颜色改为红色
                     *      3. 把祖父视为 node, 再去判断是否违反规则了
                     */
                    parent.color = uncle.color = BLACK;
                    grandpa.color = RED;
                    node = grandpa;
                    continue;
                } else {
                    // 判断 grandpa <-> parent 和 parent <-> node 的关系是否不一致
                    // 已知 parent 是 grandpa 的右边
                    if (node == parent.left) {
                        rightRotate(parent);

                        // swap(parent,node);
                        parent = node;
                    }
                    // 接下来统一处理关系一致的情况
                    leftRotate(grandpa);
                    grandpa.color = RED;
                    parent.color = BLACK;
                    break;
                }
            }
        }
        /**
         * 无论之前是什么情况，统一把根改成黑色
         * 走到此处时，root 一定不是 null
         */
        root.color = BLACK;
    }

    // 以 m 为结点，进行旋转
    private void leftRotate(RBTreeNode m) {
        // m 代表图中的 b 结点
        // parent 代表 b 结点可能存在的父亲
        RBTreeNode parent = m.parent;
        // right 代表图中的 a 结点
        RBTreeNode right = m.right;
        // leftOfRight 代表图中的可能存在的乙子树的根结点
        RBTreeNode leftOfRight = right.left;
        /*
        其中: m != null && right != null
        但是: parent 不保证 !null, leftOfRight 不保证 !null
         */

        right.parent = parent;  // 蓝色线的关系
        // 黑色线的关系
        if (parent == null) {
            // m 是 root
            root = right;
        } else {
            if (m == parent.left) {
                parent.left = right;
            } else {
                parent.right = right;
            }
        }

        right.left = m; // 黑色线的关系
        m.parent = right;   // 蓝色线的关系

        m.right = leftOfRight;
        if (leftOfRight != null) {
            leftOfRight.parent = m;
        }
    }

    private void rightRotate(RBTreeNode m) {
        RBTreeNode parent = m.parent;
        RBTreeNode left = m.left;
        RBTreeNode rightOfLeft = left.right;

        left.parent = parent;
        if (parent == null) {
            root = left;
        } else {
            if (m == parent.left) {
                parent.left = left;
            } else {
                parent.right = left;
            }
        }

        left.right = m;
        m.parent = left;

        m.left = rightOfLeft;
        if (rightOfLeft != null) {
            rightOfLeft.parent = m;
        }
    }
}

